import os

import nibabel as nib
import numpy as np
import SimpleITK as sitk
from scipy.ndimage import zoom
from skimage import exposure


class DataPreprocessor:
    """数据预处理类"""

    def __init__(self, target_spacing=(1.0, 1.0, 1.0), target_size=(256, 256)):
        self.target_spacing = target_spacing
        self.target_size = target_size

    def load_nifti(self, filepath):
        """加载NIfTI文件"""
        img = nib.load(filepath)
        data = img.get_fdata()
        affine = img.affine
        return data, affine

    def save_nifti(self, data, affine, filepath):
        """保存NIfTI文件"""
        img = nib.Nifti1Image(data, affine)
        nib.save(img, filepath)

    def resample_volume(self, volume, original_spacing, target_spacing):
        """重采样3D体积到目标间隔"""
        resize_factor = np.array(original_spacing) / np.array(target_spacing)
        new_shape = np.round(volume.shape * resize_factor).astype(int)
        resampled = zoom(volume, resize_factor, order=1)
        return resampled

    def normalize_intensity(self, volume):
        """强度归一化"""
        # 去除异常值
        p1, p99 = np.percentile(volume, (1, 99))
        volume = np.clip(volume, p1, p99)

        # 归一化到[0, 1]
        volume = (volume - np.min(volume)) / (np.max(volume) - np.min(volume))
        return volume

    def histogram_equalization(self, volume):
        """直方图均衡化"""
        # 对每个切片进行直方图均衡化
        for i in range(volume.shape[2]):
            if np.max(volume[:, :, i]) > 0:
                volume[:, :, i] = exposure.equalize_hist(volume[:, :, i])
        return volume

    def extract_2d_slices(self, volume, mask=None, min_ventricle_ratio=0.01):
        """
        从3D体积中提取2D切片

        Args:
            volume: 3D图像数据
            mask: 3D标注数据
            min_ventricle_ratio: 最小脑室占比阈值

        Returns:
            slices: 2D切片列表
            slice_masks: 对应的标注切片列表
            valid_indices: 有效切片的索引
        """
        slices = []
        slice_masks = []
        valid_indices = []

        for i in range(volume.shape[2]):
            slice_2d = volume[:, :, i]

            if mask is not None:
                mask_2d = mask[:, :, i]
                # 只保留包含足够脑室区域的切片
                ventricle_ratio = np.sum(mask_2d > 0) / mask_2d.size

                if ventricle_ratio >= min_ventricle_ratio:
                    slices.append(slice_2d)
                    slice_masks.append(mask_2d)
                    valid_indices.append(i)
            else:
                # 如果没有mask，保留所有非空切片
                if np.sum(slice_2d) > 0:
                    slices.append(slice_2d)
                    valid_indices.append(i)

        return slices, slice_masks, valid_indices

    def resize_slice(self, slice_2d, target_size):
        """调整2D切片大小"""
        if slice_2d.shape != target_size:
            zoom_factors = [t / s for t, s in zip(target_size, slice_2d.shape)]
            slice_2d = zoom(slice_2d, zoom_factors, order=1)
        return slice_2d

    def preprocess_volume(self, volume_path, mask_path=None):
        """
        完整的预处理流程

        Args:
            volume_path: 原始图像路径
            mask_path: 标注路径（可选）

        Returns:
            processed_slices: 预处理后的2D切片
            processed_masks: 预处理后的标注切片
            metadata: 元数据信息
        """
        # 加载数据
        volume, affine = self.load_nifti(volume_path)
        mask = None
        if mask_path:
            mask, _ = self.load_nifti(mask_path)

        # 强度归一化
        volume = self.normalize_intensity(volume)

        # 直方图均衡化
        volume = self.histogram_equalization(volume)

        # 提取2D切片
        slices, slice_masks, valid_indices = self.extract_2d_slices(volume, mask)

        # 调整切片大小
        processed_slices = []
        processed_masks = []

        for i, slice_2d in enumerate(slices):
            # 调整大小
            resized_slice = self.resize_slice(slice_2d, self.target_size)
            processed_slices.append(resized_slice)

            if mask is not None and i < len(slice_masks):
                resized_mask = self.resize_slice(slice_masks[i], self.target_size)
                # 确保mask是二值的
                resized_mask = (resized_mask > 0.5).astype(np.uint8)
                processed_masks.append(resized_mask)

        metadata = {
            "original_shape": volume.shape,
            "affine": affine,
            "valid_indices": valid_indices,
            "num_slices": len(processed_slices),
        }

        return np.array(processed_slices), np.array(processed_masks), metadata

    def reconstruct_3d_from_slices(self, predicted_slices, metadata):
        """从2D切片重建3D体积"""
        original_shape = metadata["original_shape"]
        valid_indices = metadata["valid_indices"]

        # 创建空的3D体积
        reconstructed = np.zeros(original_shape)

        # 将预测的切片放回原位置
        for i, idx in enumerate(valid_indices):
            if i < len(predicted_slices):
                # 调整切片大小回原始尺寸
                original_size = (original_shape[0], original_shape[1])
                resized_slice = self.resize_slice(predicted_slices[i], original_size)
                reconstructed[:, :, idx] = resized_slice

        return reconstructed


class DataAugmentation:
    """数据增强类"""

    @staticmethod
    def random_rotation(image, mask, max_angle=15):
        """随机旋转"""
        angle = np.random.uniform(-max_angle, max_angle)
        # 使用scipy或其他库实现旋转
        # 这里简化处理
        return image, mask

    @staticmethod
    def random_flip(image, mask, flip_prob=0.5):
        """随机翻转"""
        if np.random.random() < flip_prob:
            image = np.flip(image, axis=1)
            mask = np.flip(mask, axis=1)
        return image, mask

    @staticmethod
    def random_intensity_shift(image, shift_range=0.1):
        """随机强度偏移"""
        shift = np.random.uniform(-shift_range, shift_range)
        image = image + shift
        image = np.clip(image, 0, 1)
        return image